#ifndef IO_ENGINE_H
#define IO_ENGINE_H

#include <iostream>
#include <fcntl.h>
#include <unistd.h>
#include <vector>
#include <unordered_map>
#include <unordered_set>
#include <mutex>
#include <queue>
#include <memory>
#include <algorithm>
#include <atomic>
#include <cstddef>
#include <libaio.h>
#include <cstdint>
#include <condition_variable>

#define ANN_BLOCK_SIZE 4096
#define MAX_BATCH_SIZE 64 //IO引擎最多允许一次最多发起多少个读请求，大于这个值将被减小到MAX_BATCH_SIZE

#define SMALL_IO_THRESHOLD 4096
std::atomic<long long> global_io_budget(1024LL * 768);

// 为等待/预占使用的互斥与条件变量（使用 mutex+condvar 简化等待逻辑，避免复杂 CAS）
static std::mutex budget_mutex;
static std::condition_variable budget_cv;

// IO 请求封装，便于在完成时恢复预算并释放
struct IOReq {
    struct iocb cb;
    size_t size;
    int id;
    IOReq(size_t s, int _id) : size(s), id(_id) {
        memset(&cb, 0, sizeof(cb));
    }
};

// helper: 从 iocb* 得到对应的 IOReq*
static inline IOReq* req_from_iocb(struct iocb* p) {
    return reinterpret_cast<IOReq*>(
        reinterpret_cast<char*>(p) - offsetof(IOReq, cb)
    );
}

class IOUringManager_ASYNC {
private:
    int fd_;
    int numRequests_;
    io_context_t ctx_;

public:
    IOUringManager_ASYNC(int fd, int request_count)
        : fd_(fd), numRequests_(request_count) {
        memset(&ctx_, 0, sizeof(io_context_t));
        int ret = io_setup(numRequests_, &ctx_);
        if (ret < 0) {
            std::cerr << "Failed to initialize libaio. Error code: " << ret << std::endl;
            exit(EXIT_FAILURE);
        }
    }

    ~IOUringManager_ASYNC() {
        io_destroy(ctx_);
    }

private:
    // 阻塞直到可以为 req_size 预占预算，然后返回（已从 global_io_budget 中扣减）
    void wait_and_reserve_budget(long long req_size) {
        // 小 IO 不消耗预算（或者消耗很少）
        if (req_size <= SMALL_IO_THRESHOLD) {
            return; 
        }

        std::unique_lock<std::mutex> lock(budget_mutex);
        budget_cv.wait(lock, [&] {
            return global_io_budget.load() >= req_size;
        });
        global_io_budget.fetch_sub(req_size);
    }

    // 释放被占用的预算并唤醒等待者
    void release_budget_and_notify(long long bytes) {
        if (bytes <= 0) return;
        global_io_budget.fetch_add(bytes, std::memory_order_relaxed);
        budget_cv.notify_all();
    }

public:
    // ids: 用户请求 id 向量。input_num_request 可用于只提交前 N 个请求
    // 语义：要么成功提交（返回提交数量），要么阻塞直到提交成功或出现不可恢复错误（此实现倾向于重试），不会简单返回失败让上层退出
    int submitRequests(const std::vector<void*>& input_bufAddrs,
                       const std::vector<long long>& offsets,
                       const std::vector<long long>& sizes,
                       const std::vector<int>& ids,
                       int input_num_request = 0) {

        size_t num_requests = input_num_request == 0 ? input_bufAddrs.size() : input_num_request;
        if (num_requests == 0) return 0;
        if (input_bufAddrs.size() < num_requests || offsets.size() < num_requests ||
            sizes.size() < num_requests || ids.size() < num_requests) {
            std::cerr << "Input vectors too small for requested operations" << std::endl;
            return 0;
        }

        // 先为每个请求按顺序阻塞预占预算（避免上层需要处理部分提交失败）
        // 注意：按请求顺序预占会自然让小请求（size 小）先拿到预算（因为 req_size 小更容易满足）
        std::vector<IOReq*> allocated_reqs;
        allocated_reqs.reserve(num_requests);
        std::vector<struct iocb*> iocb_ptrs;
        iocb_ptrs.reserve(num_requests);

        try {
            for (size_t i = 0; i < num_requests; ++i) {
                size_t req_size = static_cast<size_t>(sizes[i]);
                // 等待并预占预算（严格阻塞直到成功）
                wait_and_reserve_budget(req_size);

                // 预占成功 -> 分配请求结构
                IOReq* req = new IOReq(req_size, ids[i]);
                io_prep_pread(&req->cb, fd_, input_bufAddrs[i], req_size, offsets[i]);
                req->cb.data = reinterpret_cast<void*>(static_cast<uintptr_t>(ids[i]));
                allocated_reqs.push_back(req);
                iocb_ptrs.push_back(&req->cb);
            }
        } catch (...) {
            // 出现异常：恢复已预占预算并释放内存
            long long restored = 0;
            for (IOReq* r : allocated_reqs) {
                restored += static_cast<long long>(r->size);
                delete r;
            }
            if (restored > 0) release_budget_and_notify(restored);
            return 0;
        }

        if (iocb_ptrs.empty()) return 0;

        // 现在所有请求都已预占预算并准备好，开始提交
        int to_submit = static_cast<int>(iocb_ptrs.size());
        int submitted = 0;

        // 我们采用循环提交的方式：如果 io_submit 只提交了部分，我们提交剩余的。
        // 在不可恢复错误时，恢复未提交的预算并返回 0（或根据需要可以选择退出程序）
        while (submitted < to_submit) {
            int remain = to_submit - submitted;
            int ret = io_submit(ctx_, remain, iocb_ptrs.data() + submitted);
            if (ret >= 0) {
                submitted += ret;
                // 若全部提交完毕，跳出循环
            } else {
                // ret < 0 表示出错，尝试处理常见可重试错误
                int err = -ret;
                // 对于 EINTR/EAGAIN 之类的，短等待后重试
                if (err == EINTR || err == EAGAIN) {
                    std::this_thread::sleep_for(std::chrono::milliseconds(1));
                    continue;
                } else {
                    // 不可恢复错误：恢复尚未提交请求的预算并释放内存
                    long long restore_bytes = 0;
                    for (int i = submitted; i < to_submit; ++i) {
                        IOReq* r = req_from_iocb(iocb_ptrs[i]);
                        restore_bytes += static_cast<long long>(r->size);
                    }
                    release_budget_and_notify(restore_bytes);

                    // 释放所有已分配对象（包括已提交的那些，已提交的会在完成时再次 delete -> 导致 double free）
                    // **因此在不可恢复错误情况下，需要谨慎处理**。这里假定不可恢复错误很少发生，并把所有 allocated 释放掉。
                    for (IOReq* r : allocated_reqs) delete r;
                    std::cerr << "io_submit fatal error: " << err << std::endl;
                    return 0;
                }
            }
        }

        // 成功把所有 iocb 提交到内核，返回提交数量
        return submitted;
    }

    // 获取完成的请求 id 列表，并恢复预算（唤醒等待者）
    int getCompletedResultIds(std::vector<int>& completedIds, int min_nr = 0) {
        struct io_event events[MAX_BATCH_SIZE];
        int ret = io_getevents(ctx_, min_nr, MAX_BATCH_SIZE, events, NULL);
        if (ret < 0) {
            std::cerr << "Failed to wait for I/O completion" << std::endl;
            exit(EXIT_FAILURE);
        }

        long long restored_bytes = 0;
        for (int i = 0; i < ret; ++i) {
            struct iocb* cb = static_cast<struct iocb*>(events[i].obj);
            IOReq* r = req_from_iocb(cb);
            int user_id = static_cast<int>(reinterpret_cast<uintptr_t>(cb->data));

            if (events[i].res < 0) {
                std::cerr << "I/O operation failed for request " << user_id
                          << " err=" << events[i].res << std::endl;
                // 失败也应该恢复预算（因为之前为该请求预占了）
            } else {
                completedIds.push_back(user_id);
            }

            restored_bytes += static_cast<long long>(r->size);
            delete r;
        }

        if (restored_bytes > 0) {
            release_budget_and_notify(restored_bytes);
        }
        return ret;
    }
};


//一些便于Direct IO的帮助函数
//计算对齐的块数
int block_count(long long size) {
    return (size + ANN_BLOCK_SIZE - 1) / ANN_BLOCK_SIZE;
}

//计算对齐的大小。输入需要的字节数，返回大于等于需要的字节数的对齐大小.
long long aligned_size(long long size, long long start_addr = 0) {
    return block_count(size + start_addr%ANN_BLOCK_SIZE) * ANN_BLOCK_SIZE;
}

//计算对齐的偏移量，即某个偏移量在对齐IO中，在文件中的起始位置，≤起始地址
long long aligned_offset_in_file(long long offset) {
    return offset / ANN_BLOCK_SIZE * ANN_BLOCK_SIZE;
}

//通过aligned_offset_in_file返回的地址中，真实起始地址在分配的buf中的偏移量
long long aligned_offset_in_buf(long long offset) {
    return offset % ANN_BLOCK_SIZE;
}

//给定需要的字节数，返回≥字节数且对齐的内存地址
void* get_aligned_buffer(long long size, long long start_addr = 0) {
    if (size <= 0) {
        std::cout << "[Error] allocate_aligned_buffer size is invalid" << std::endl;
        exit(0);
    }
    // 分配对齐的内存
    void* buf = nullptr;
    if (posix_memalign(&buf, ANN_BLOCK_SIZE, aligned_size(size, start_addr)) != 0) {
        std::cout << "[Error] posix_memalign failed" << std::endl;
        exit(0);
    }
    return buf;
}

//用O_DIRECT方式打开文件后用该方法读取文件
//负责对buf分配空间，并且使其指向真实数据的地址
//现在测试阶段为了方便使用了拷贝，后续进行优化
int read_file_aligned(int fd, char* buf, long long offset, long long size) {
    if(size <= 0) {
        exit(0); // 无效的大小
    }
    // 计算需要的对齐内存大小
    int aligned_size = (size + ANN_BLOCK_SIZE - 1) / ANN_BLOCK_SIZE * ANN_BLOCK_SIZE;

    // 分配对齐的内存
    void* aligned_buf = nullptr;
    if (posix_memalign(&aligned_buf, ANN_BLOCK_SIZE, aligned_size) != 0) {
        perror("posix_memalign");
        exit(0);
    }

    // 设置文件偏移量
    int aligned_offset = offset / ANN_BLOCK_SIZE * ANN_BLOCK_SIZE;
    if (lseek(fd, aligned_offset, SEEK_SET) == -1) {
        perror("lseek");
        free(aligned_buf);
        exit(0);
    }

    // 读取文件内容到对齐的内存中
    ssize_t bytes_read = read(fd, aligned_buf, aligned_size);
    if (bytes_read == -1 || bytes_read < size) {
        perror("read");
        free(aligned_buf);
        exit(0);
    }

    //计算offset在这个buf中的偏移量
    int offset_in_buf = offset % ANN_BLOCK_SIZE;
    // 将对齐的内存内容复制到buf中
    memcpy(buf, (char*)aligned_buf + offset_in_buf, size);
    return bytes_read;
}

#endif